import gzip
import os
import pickle
import sys
import numpy as np
from network import NeuralNetwork
def load_mnist():
    path = '../mnist_data/data/mnist.pkl.gz'  # 下载的mnist数据存放路径,
    data_file = gzip.open(path, "rb")  # 读取压缩文件
    train_data, val_data, test_data = pickle.load(data_file, encoding="latin1")  # 将数据解压缩 方式为latin1
    data_file.close()  # 关闭文件
    train_inputs = [np.reshape(x, (784, 1)) for x in train_data[0]]  # 训练集输入
    train_results = [vectorized_result(y) for y in train_data[1]] #  对训练集结果进行向量化
    train_data = list(zip(train_inputs, train_results))  # 组合训练集输入和结果

    val_inputs = [np.reshape(x, (784, 1)) for x in val_data[0]] #将每个输入数据 `x` 重塑为形状为 `(784, 1)` 的数组。MNIST 图像的大小是 `28x28`，所以 `784` 是 `28x28` 的展平形式。
    val_results = val_data[1] #获取标签数据
    val_data = list(zip(val_inputs, val_results))

    test_inputs = [np.reshape(x, (784, 1)) for x in test_data[0]]
    test_data = list(zip(test_inputs, test_data[1]))
    return train_data, val_data, test_data

import matplotlib.pyplot as plt
def display_and_save_images(test_data, predictions, num_images=10, filename='predictions.png'):
    fig, axes = plt.subplots(1, num_images, figsize=(15, 15))  # 创建一个10x10的子图
    for i in range(num_images):  # 遍历每张图片
        ax = axes[i]  # 获取子图
        image = test_data[i][0].reshape(28, 28)  # 重塑图片为28x28
        ax.imshow(image, cmap='gray')  # 显示灰度图
        ax.set_title(f"Pred: {predictions[i]}")  # 显示预测结果
        ax.axis('off')
    plt.savefig(filename)
    plt.show()


def vectorized_result(y):  # 定义向量化结果
    e = np.zeros((10, 1))  # 初始化一个10x1的数组
    e[y] = 1.0  # 将y位置的值设为1
    return e
#定义保存模型的方法

def save_predictions(predictions, filename='predictions.pkl'):
    with open(filename, 'wb') as f:
        pickle.dump(predictions, f)  # 保存预测结果


if __name__ == "__main__":
    np.random.seed(42)  # 设置随机种子
    layers = [784, 30, 10]  # 网络结构
    learning_rate = 0.01  # 学习率
    mini_batch_size = 16  # 批处理大小
    epochs = 1  # 训练轮数
    #如果按下y 进入训练模式, 否则进入测试模式
    if input("按下y 进入训练模式, 否则进入测试模式") == "y":
        # 初始化训练、val 和测试数据
        train_data, val_data, test_data = load_mnist()  # 加载数据
        nn = NeuralNetwork(layers, learning_rate, mini_batch_size, "relu")  # 初始化神经网络
        nn.fit(train_data, val_data, epochs)  # 训练神经网络
        # 测试神经网络
        print("Testing the network...")
        accuracy = nn.validate(test_data) / 100.0  # 计算测试集准确率
        print(f"Test Accuracy: {accuracy}%.")  # 打印测试集准确率
        nn.save()
        #保存训练结果
        # 获取预测结果
        predictions = [nn.predict(x) for x, _ in test_data]
        # 显示并保存测试结果图像
        display_and_save_images(test_data, predictions, num_images=10, filename='test_predictions.png')
    else:
        # 加载模型并测试
        if not os.path.exists('models/model.npz'):
            print("模型不存在")
            sys.exit(1)
        nn = NeuralNetwork(layers, learning_rate, mini_batch_size, "relu")  # 初始化神经网
        nn.load()  # 络
          # 测试神经网络
        print("Testing the network...")
        train_data, val_data, test_data = load_mnist()
        accuracy = nn.validate(test_data) / 100.0  # 计算测试集准确率
        print(f"Test Accuracy: {accuracy}%.")  # 打印测试集准确率
        # 获取预测结果
        predictions = [nn.predict(x) for x, _ in test_data]
        # 保存预测结果
        #save_predictions(predictions, filename='predictions.pkl')
        # 显示并保存测试结果图像
        display_and_save_images(test_data, predictions, num_images=10, filename='test_predictions.png')
